package com.tao.auth.handler;

import com.fasterxml.jackson.databind.ObjectMapper;
import jakarta.servlet.ServletException;
import jakarta.servlet.http.HttpServletRequest;
import jakarta.servlet.http.HttpServletResponse;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.http.MediaType;
import org.springframework.security.core.Authentication;
import org.springframework.security.core.AuthenticationException;
import org.springframework.security.oauth2.core.*;
import org.springframework.security.oauth2.core.endpoint.OAuth2AuthorizationRequest;
import org.springframework.security.oauth2.jose.jws.SignatureAlgorithm;
import org.springframework.security.oauth2.jwt.JwtClaimsSet;
import org.springframework.security.oauth2.server.authorization.OAuth2Authorization;
import org.springframework.security.oauth2.server.authorization.OAuth2AuthorizationService;
import org.springframework.security.oauth2.server.authorization.OAuth2TokenType;
import org.springframework.security.oauth2.server.authorization.client.RegisteredClient;
import org.springframework.security.oauth2.server.authorization.client.RegisteredClientRepository;
import org.springframework.security.oauth2.server.authorization.context.AuthorizationServerContext;
import org.springframework.security.oauth2.server.authorization.context.AuthorizationServerContextHolder;
import org.springframework.security.oauth2.server.authorization.settings.AuthorizationServerSettings;
import org.springframework.security.oauth2.server.authorization.token.*;
import org.springframework.security.web.authentication.AuthenticationSuccessHandler;
import org.springframework.stereotype.Component;
import org.springframework.security.core.GrantedAuthority;
import org.springframework.security.oauth2.jwt.*;

import java.security.Principal;
import java.time.Duration;
import java.time.Instant;
import java.time.temporal.ChronoUnit;
import java.io.IOException;
import java.util.*;
import java.util.stream.Collectors;

/**
 * <p>
 * 自定义成功处理器
 * </p>
 *
 * @author taohongrun
 * @since 2024/11/23
 */

@Component
public class CustomAuthenticationSuccessHandler implements AuthenticationSuccessHandler {

    @Autowired
    private OAuth2TokenGenerator<?> tokenGenerator;

    @Autowired
    private RegisteredClientRepository registeredClientRepository;

    @Autowired
    private AuthorizationServerSettings authorizationServerSettings;

    @Autowired
    private OAuth2AuthorizationService authorizationService;

    @Override
    public void onAuthenticationSuccess(HttpServletRequest request, HttpServletResponse response,
                                        Authentication authentication) throws IOException, ServletException {
        try {
            RegisteredClient registeredClient = registeredClientRepository.findByClientId("LearnHubWebApp");
            if (registeredClient == null) {
                throw new IllegalStateException("RegisteredClient not found");
            }

            Set<String> scopes = new HashSet<>(registeredClient.getScopes());

            // 创建授权构建器
            OAuth2Authorization.Builder authorizationBuilder = OAuth2Authorization.withRegisteredClient(registeredClient)
                    .principalName(authentication.getName())
                    .id(UUID.randomUUID().toString())
                    .authorizationGrantType(AuthorizationGrantType.PASSWORD)
                    .authorizedScopes(scopes)
                    // 添加认证信息作为属性，这样在 token 生成时可以获取到
                    .attribute(Principal.class.getName(), authentication);

            // 设置授权服务器上下文
            AuthorizationServerContext context = new AuthorizationServerContext() {
                @Override
                public String getIssuer() {
                    return authorizationServerSettings.getIssuer();
                }

                @Override
                public AuthorizationServerSettings getAuthorizationServerSettings() {
                    return authorizationServerSettings;
                }
            };

            AuthorizationServerContextHolder.setContext(context);

            // 构建 token 上下文
            DefaultOAuth2TokenContext.Builder tokenContextBuilder = DefaultOAuth2TokenContext.builder()
                    .registeredClient(registeredClient)
                    .principal(authentication)
                    .authorizationServerContext(context)
                    .authorization(authorizationBuilder.build())  // 添加授权信息
                    .authorizedScopes(scopes)
                    .tokenType(OAuth2TokenType.ACCESS_TOKEN)
                    .authorizationGrantType(AuthorizationGrantType.PASSWORD);

            // 生成令牌
            OAuth2Token generatedToken = tokenGenerator.generate(tokenContextBuilder.build());
            if (generatedToken == null) {
                throw new IllegalStateException("Generated token is null");
            }

            if (generatedToken instanceof Jwt) {
                Jwt jwt = (Jwt) generatedToken;
                OAuth2AccessToken accessToken = new OAuth2AccessToken(
                        OAuth2AccessToken.TokenType.BEARER,
                        jwt.getTokenValue(),
                        jwt.getIssuedAt(),
                        jwt.getExpiresAt(),
                        scopes
                );

                // 构建授权对象
                OAuth2Authorization authorization = authorizationBuilder
                        .token(accessToken, metadata -> {
                            metadata.put(OAuth2Authorization.Token.CLAIMS_METADATA_NAME, jwt.getClaims());
                        })
                        .build();

                // 保存授权信息
                authorizationService.save(authorization);

                // 构建响应
                Map<String, Object> responseMap = new HashMap<>();
                responseMap.put("access_token", jwt.getTokenValue());
                responseMap.put("token_type", "Bearer");
                responseMap.put("expires_in", getExpiresIn(jwt));
                // 从 JWT claims 中获取权限信息
                @SuppressWarnings("unchecked")
                List<String> authorities = (List<String>) jwt.getClaims().get("authorities");
                responseMap.put("authorities", authorities);

                response.setContentType(MediaType.APPLICATION_JSON_VALUE);
                response.setCharacterEncoding("UTF-8");
                new ObjectMapper().writeValue(response.getWriter(), responseMap);
            }
        } catch (Exception e) {
            throw new AuthenticationException("Token generation failed", e) {};
        } finally {
            AuthorizationServerContextHolder.resetContext();
        }
    }

    // 计算token过期时间
    private long getExpiresIn(Jwt jwt) {
        if (jwt.getExpiresAt() != null) {
            return Duration.between(Instant.now(), jwt.getExpiresAt()).getSeconds();
        }
        return 7200; // 默认2小时
    }

}


